#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @Time    : 2022/3/20 1:22 下午
# @Author  : WangZhixing
# @Author  : WangZhixing


from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from torch_geometric.utils import train_test_split_edges

from Metric import Cluster_Metric
from ..Module.Encoder import Encoder3, VEncoder1
from ..Module.GAE import GAE
import torch

def unsurpervised_train(model, optimizer, data):
    model.train()
    optimizer.zero_grad()
    z = model.encode(data.x, data.train_pos_edge_index)
    loss = model.recon_loss(z, data.train_pos_edge_index)
    loss.backward()
    optimizer.step()
    return loss

def unsurpervised_test(model, data):
    model.eval()
    with torch.no_grad():
        z = model.encode(data.x, data.train_pos_edge_index)
    kmeans_input = z.cpu().numpy()
    kmeans = KMeans(n_clusters=14, random_state=0).fit(kmeans_input)
    pred = kmeans.predict(kmeans_input)
    labels = data.y.cpu().numpy()
    com, hm, nmi = Cluster_Metric(labels, pred)
    # 前两个数的平均,相当于聚类的精确度。
    roc_auc_score, average_precision_score = model.test(z, data.test_pos_edge_index, data.test_neg_edge_index)
    return roc_auc_score, average_precision_score, com, hm, nmi


def GAE_train(data, **kwarg):
    encoder = Encoder3(data.num_features, 28)
    model = GAE(encoder)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=0)
    data = train_test_split_edges(data)
    for epoch in range(kwarg['train_epoch']):
        loss = unsurpervised_train(model, optimizer, data)
    with torch.no_grad():
        z = model.encode(data.x, data.train_pos_edge_index)
    kmeans_input = z.cpu().numpy()
    kmeans = KMeans(n_clusters=kwarg['cluster'], random_state=0).fit(kmeans_input)
    preds = kmeans.predict(kmeans_input)
    return kmeans_input, preds
